package com.leetcode.p300;

import com.utils.DataUtils;

import java.util.Comparator;
import java.util.HashMap;
import java.util.PriorityQueue;
import java.util.function.BiConsumer;

public class Test347 {

    public int[] topKFrequent(int[] nums, int k) {
        HashMap<Integer, Integer> hashMap = new HashMap<>();
        for (int num : nums) {
            if (hashMap.containsKey(num)) {
                hashMap.put(num, hashMap.get(num) + 1);
            } else {
                hashMap.put(num, 1);
            }
        }

        PriorityQueue<int[]> minHeap = new PriorityQueue<>(11, new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return o1[1] - o2[1];
            }
        });

        hashMap.forEach(new BiConsumer<Integer, Integer>() {
            @Override
            public void accept(Integer key, Integer value) {
                if (minHeap.size() == k) {
                    if (minHeap.peek()[1] < value) {
                        minHeap.poll();
                        minHeap.offer(new int[]{key, value});
                    }
                } else {
                    minHeap.offer(new int[]{key, value});
                }

            }
        });

        int[] res = new int[k];

        for (int i = 0; i < k; i++) {
            res[i] = minHeap.poll()[0];
        }

        return res;
    }

    public static void main(String[] args) {
        Test347 test = new Test347();
        DataUtils.printArray(test.topKFrequent(new int[]{1, 1, 1, 2, 2, 3}, 2));
    }
}
